"""
    @brief:simulation environment
    @Editor:CJH
    @Date:2025/4/24
"""
import sys
import os.path
import threading
import time
sys.path.append("utils")

import cv2
import numpy as np
from dataloader import Dataloader

font = cv2.FONT_HERSHEY_COMPLEX_SMALL
import matplotlib.pyplot as plt

start = time.time()


# 机器人对象
class BOT:
    def __init__(self, name: str, start_point: tuple, total_stride: int):
        self.name = name
        self.start_point = start_point  # np.array(x,y)
        self.total_stride = total_stride


class Colors:
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hexs = ('344593', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
                '2C99A8', '00C2FF', 'FF3838', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        self.n = len(self.palette)

    def rgb_call(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))

    @staticmethod
    def rgb2hex(rgb):
        """RGB转HEX
        :param rgb: RGB颜色元组，Tuple[int, int, int]
        :return: int or str
        """
        r, g, b = rgb
        result = r + (g << 8) + (b << 16)
        return hex(result)[2:]


# 只作为构建一个动态地图环境，可随时更新，用于训练DQN模型
class Map():
    def __init__(self, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20):  # split_nums将地图分割为不同大小得栅格地图
        super(Map, self).__init__()
        # 定义地图相关信息
        self.title = "Rectangle board"
        self.canvas_shape = (*canvas_size, 3)  # shape=(y,x)
        self.grid_map = np.zeros((split_y_nums, split_x_nums))  # shape=(y,x)
        self.mass_fraction_viewer = Dataloader(self.grid_map.shape)
        # Declare all maps
        self.canvas, self.mass_fraction_map, self.Gaussian_mass_map = None, None, None
        # Define elements present inside the environment
        # 构造地图边界，确保机器人不会越界
        self.y_min = 0
        self.x_min = 0
        self.y_max = self.canvas_shape[0]
        self.x_max = self.canvas_shape[1]

    # 可视化地图初始化
    def map_init(self, csv_path="", update_=False):
        canvas = np.ones(self.canvas_shape, dtype=np.uint8) * 255
        # draw row line
        for row in range(1, self.grid_map.shape[0] + 1):
            cv2.line(canvas, (0, int(self.canvas_shape[0] * row / self.grid_map.shape[0])),
                     (self.canvas_shape[1], int(self.canvas_shape[0] * row / self.grid_map.shape[0])), (255, 0, 0), 1)

        for col in range(1, self.grid_map.shape[1] + 1):
            cv2.line(canvas, (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), 0),
                     (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), self.canvas_shape[0]), (0, 0, 255), 1)

        # 更新气体质量分数可视化地图
        if os.path.exists(csv_path) and update_:
            # 载入真实值的气体分数地图
            self.mass_fraction_viewer.merge(csv_path)
            self.mass_fraction_map = self.mass_fraction_viewer.init_map
            # 载入模拟真实情况的高斯噪声地图
            self.Gaussian_mass_map = self.mass_fraction_viewer.noisy_map
            # 融合气体质量分数地图与栅格地图
            alpha = 0.3  # 第一幅图像的权重
            beta = 0.7  # 第二幅图像的权重
            gamma = 0.5  # 亮度调整
            mass_fraction_map = cv2.resize(self.mass_fraction_map, self.canvas_shape[1::-1])
            depth_image = cv2.normalize(mass_fraction_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)
            canvas = cv2.addWeighted(canvas, alpha, depth_image, beta,
                                     gamma)
        self.canvas = canvas

    # 多线程展现画布
    def show(self):
        while True:
            cv2.imshow(self.title, self.canvas)
            cv2.waitKey(10)


class Point(Colors):
    def __init__(self, index, name, x_max, x_min, y_max, y_min, total_stride, icon_shape=(50, 50)):
        super(Point, self).__init__()
        self.x = 0
        self.y = 0
        self.x_min = x_min
        self.x_max = x_max
        self.y_min = y_min
        self.y_max = y_max
        self.index = index
        self.name = name
        # 记录开始位置
        self.start_point = None
        self.current_grid_pos = None  # x,y
        self.current_pixel_pos = None  # x,y
        self.icon = np.full((*icon_shape, 3), self.rgb_call(index), dtype=np.uint8)
        self.pixel_list = []  # 记录机器走过的路径，像素坐标(x,y)
        self.grid_pos_list = []  # 记录机器走过的路径，栅格坐标(x,y)
        # 初始总步长
        self.total_stride = total_stride
        self.ep_return = 0  # 回归经验值
        self.stride_left = self.total_stride  # 训练步长

    def set_position(self, x, y, start_point):
        self.x = self.clamp(x, self.x_min, self.x_max - self.icon.shape[1])
        self.y = self.clamp(y, self.y_min, self.y_max - self.icon.shape[0])
        # 初始化，添加当前位置
        self.current_pixel_pos = (self.x, self.y)
        self.current_grid_pos = start_point
        # 初始化，添加当前位置
        self.pixel_list.append((self.x, self.y))
        self.grid_pos_list.append(self.current_grid_pos)

    def move(self, del_x, del_y, sim_):
        next_x = del_x + self.x
        next_y = del_y + self.y
        # 如果移动之后的距离超出了，返回False表明这次的移动是失败的
        if self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - self.icon.shape[
            0] and not sim_:
            self.x += del_x
            self.y += del_y
            # 更新当前坐标
            self.current_pixel_pos = (self.x, self.y)
            return True
        elif self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - \
                self.icon.shape[
                    0] and sim_:
            return True
        else:
            return False

    def clamp(self, n, minn, maxn):
        return max(min(maxn, n), minn)


# 终端用于控制每个机器人，便于生成多个机器人
class Terminal(Map):
    def __init__(self, bot_elements, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20):
        Map.__init__(self, canvas_size=canvas_size, split_x_nums=split_x_nums, split_y_nums=split_y_nums)
        # 初始化地图
        self.map_init()
        # 通过一个字典管理所有机器人
        self.element_collection = self.create_bot(bot_elements)
        # 卷积核，使机器人移动一步的时候，有周围一圈的结果
        self.kernel = np.ones((3, 3), np.uint8)
        # 离散化移动的方式以及步骤
        self.box_min_step = np.array([int(self.canvas_shape[1] / self.grid_map.shape[1]),
                                      int(self.canvas_shape[0] / self.grid_map.shape[0])])
        self.action_dict = {"upper": self.box_min_step * np.array([0, -1]),
                            "lower": self.box_min_step * np.array([0, 1]),
                            "left": self.box_min_step * np.array([-1, 0]),
                            "right": self.box_min_step * np.array([1, 0]),
                            "up_right": self.box_min_step * np.array([1, -1]),
                            "up_left": self.box_min_step * np.array([-1, -1]),
                            "low_right": self.box_min_step * np.array([1, 1]),
                            "low_left": self.box_min_step * np.array([-1, 1])}
        self.known_max_pos_after = None

    # 将所有元素呈现在画布上
    def draw_elements_on_canvas(self, csv_path):
        # Init the canvas
        self.map_init(csv_path=csv_path, update_=True)
        # 将每个机器人画出来
        for index, Bot_name in enumerate(self.element_collection.keys()):
            text = f"Bot-{Bot_name} | Stride Left: {self.element_collection[Bot_name].stride_left} | Rewards: {self.element_collection[Bot_name].ep_return}"
            # Put the info on canvas
            cv2.putText(self.canvas, text, (10, 20 * (index + 1)), font,
                        1, (0, 0, 0), 1, cv2.LINE_AA)
            # Draw the heliopter on canvas
            x, y = self.element_collection[Bot_name].x, self.element_collection[Bot_name].y
            elem_shape = self.element_collection[Bot_name].icon.shape
            self.canvas[y: y + elem_shape[0], x:x + elem_shape[1]] = self.element_collection[Bot_name].icon
            # for index, point in enumerate(self.element_collection[Bot_name].pixel_list):
            #     if index < len(self.element_collection[Bot_name].pixel_list) - 1:
            #         cv2.line(self.canvas, self.element_collection[Bot_name].pixel_list[index], self.element_collection[Bot_name].pixel_list[index + 1], (0, 255, 0), 5)
            #     cv2.circle(self.canvas, point, 5, (0, 255, 0), -1)

    def create_bot(self, bot_elements: dict):
        element_collection = {}
        icon_shape = (
            int((self.canvas_shape[1] / self.grid_map.shape[1]) * 0.3),
            int((self.canvas_shape[0] / self.grid_map.shape[0]) * 0.3))

        origin_step = (int(self.canvas_shape[1] / (2 * self.grid_map.shape[1]) - icon_shape[1] / 2),
                       int(self.canvas_shape[0] / (2 * self.grid_map.shape[0]) - icon_shape[0] / 2))

        # 统一处理字典或列表
        elements = bot_elements.values() if isinstance(bot_elements, dict) else bot_elements
        for index, element in enumerate(elements):
            # 实例化机器人本体
            element_collection[element.name] = Point(index, name=element.name, x_max=self.x_max, x_min=self.x_min,
                                                     y_max=self.y_max, y_min=self.y_min,
                                                     total_stride=element.total_stride,
                                                     icon_shape=icon_shape)
            # 设定初始位置
            start_point = element.start_point
            if start_point[0] < self.grid_map.shape[0] and start_point[1] < self.grid_map.shape[0]:
                x_pos = int(start_point[0] * self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(start_point[1] * self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])
                self.grid_map[start_point[1], start_point[0]] = 1  # 更新栅格地图
            else:
                x_pos = int(self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])

            # 初始化，设置所有点坐标
            element_collection[element.name].set_position(x_pos, y_pos, start_point)

        return element_collection

    # 更新机器人在栅格地图中位置的函数
    def upgrade_grid_position(self, bot_name):
        current_state = self.element_collection[bot_name].current_pixel_pos
        current_state = int(current_state[0] * self.grid_map.shape[1] / self.canvas_shape[1]), int(
            current_state[1] * self.grid_map.shape[0] / self.canvas_shape[0])
        # 记录已经走过的路径
        self.grid_map[current_state[1], current_state[0]] = 1  # 在栅格地图上赋予机器人位置信息

        self.element_collection[bot_name].current_grid_pos = current_state  # 为机器人附上栅格地图的位置

    def take_actions(self, action, sim_, bot_name):  # 给予一个sim_仿真符号，为了测试运动是否可行的同时，不更改机器人动作
        if not isinstance(action, str):
            action = list(self.action_dict.keys())[action]
        assert self.element_collection[bot_name] is not None, "Please add element to execute following task"
        assert action in self.action_dict.keys(), "Invalid Action"
        move_flag = self.element_collection[bot_name].move(self.action_dict[action][0],
                                                           self.action_dict[action][1], sim_)

        if not sim_:
            self.upgrade_grid_position(bot_name)
            # 移动机器人，减少可移动步数，如果是假设移动，则不减少步数
            self.element_collection[bot_name].stride_left -= 1
            if move_flag:
                # 将机器人走的路径添加到坐标里
                self.element_collection[bot_name].pixel_list.append(self.element_collection[bot_name].current_pixel_pos)
                self.element_collection[bot_name].grid_pos_list.append(
                    self.element_collection[bot_name].current_grid_pos)

        # 返回布尔值，确定这一步是否是正常的
        return move_flag

    def step(self, action, bot_name, view=False):
        if not isinstance(action, str):
            action = list(self.action_dict.keys())[action]
        assert self.element_collection[bot_name] is not None, "Please add element to execute following task"
        assert action in self.action_dict.keys(), "Invalid Action"
        # parameter setting
        k1, k2, k3, k4, k5, k6, k7 = 0.5, 2, 1, 0.1, 0.5, 4, 10  # 设置权重影响分数
        stride_ratio = self.element_collection[bot_name].stride_left / self.element_collection[bot_name].total_stride
        visited_ratio = self.grid_map.sum() / (self.grid_map.shape[0] * self.grid_map.shape[1])  # 时间因素
        reward = 0
        max_mass_thresh = np.max(self.mass_fraction_map)

        done_s = False  # 定义这次step的是否完成的标志
        # 当前位置
        current_pos = self.element_collection[bot_name].current_grid_pos
        # 当前位置的气体质量分数值
        current_this_point_mass = self.mass_fraction_map[current_pos[1], current_pos[0]]
        # 移动前的位置和
        before_point_sum = self.grid_map.sum()

        if not self.take_actions(action, sim_=False, bot_name=bot_name):
            reward -= 5
            pass
            # print("撞墙啦")

        # -------------------------------------after_taking_actions-----------------------------------
        # 移动后的位置
        next_pos = self.element_collection[bot_name].current_grid_pos
        # 移动后的位置的气体质量分数值
        next_this_point_mass = self.mass_fraction_map[next_pos[1], next_pos[0]]
        # 移动后的位置和
        after_point_sum = self.grid_map.sum()
        # 1.转换为ndarray数组输出的地图，表示机器人当前位置
        pos_grid_map = np.zeros(self.grid_map.shape)
        pos_grid_map[next_pos[1], next_pos[0]] = 1
        # 2.转换为ndarray数组输出的地图，表示质量分数地图
        path_array_uint8 = (self.grid_map > 0).astype(np.uint8)
        # # 执行膨胀操作
        dilated_path = cv2.dilate(path_array_uint8, self.kernel, iterations=1)
        # 由于传感器误差，所以最后得到的气体质量分数地图应为带有高斯噪声的
        # print(dilated_path)
        # print(self.mass_fraction_map)
        quality_scores_map = np.array(dilated_path * self.mass_fraction_map)
        # # 将quality_scores_map放缩到[0,1]
        quality_scores_map = quality_scores_map / (np.max(quality_scores_map) + 1e-8)
        # # 将0变成-1，代表未知区域
        quality_scores_map = np.where(quality_scores_map == 0, -1, quality_scores_map)

        # DQN policy
        def policy_():
            # 1.Exploration Reward
            visited_reward = k3 * (1 - visited_ratio)  # 初始 5，随访问增加减小
            exploration_bonus = stride_ratio * (
                visited_reward if after_point_sum - before_point_sum != 0 else 0)  # 如果是没探索过的区域，则给予奖励，随时间变化这个影响会越来越少

            # 2.Dynamic Probability Reward 来回震荡，并且离阈值差的很远则扣分
            dynamic_Probability = -1 if after_point_sum - before_point_sum == 0 and abs(
                next_this_point_mass - max_mass_thresh) > k4 * 1.5 else 0

            # 3.New height bonus
            self.known_max_pos_before = np.unravel_index(np.argmax(quality_scores_map), quality_scores_map.shape)  # y,x
            if self.known_max_pos_after != self.known_max_pos_before:  # 只要不重复返回刷分就行
                distance_reward = np.max(quality_scores_map) / max_mass_thresh  # 如果发现一个浓度更高的地方，则奖励，归一化奖励值
                self.known_max_pos_after = self.known_max_pos_before
            else:
                distance_reward = 0

            # 4.Gradient ascent reward bonus
            Gradient_reward = (next_this_point_mass - current_this_point_mass) / max_mass_thresh
            Gradient_reward = -1 / np.exp(Gradient_reward) * k1 if Gradient_reward < 0 else np.exp(Gradient_reward)

            # 5.Closer target reward bonus
            # 如果靠得越近加的越多，离远了不加分
            fractions_max_pos_ = np.unravel_index(np.argmax(self.mass_fraction_map),
                                                  self.mass_fraction_map.shape)  # y,x
            distance_ = np.linalg.norm(np.array(next_pos) - np.array(fractions_max_pos_[::-1]))
            Closer_reward = 10 / (distance_ + 1e-6)  # 距离越近，奖励越高，如果是远离则惩罚

            # print(k5 * exploration_bonus, '探索奖励')
            # print(k6 * dynamic_Probability, '来回震荡惩罚')
            # print(k2 * distance_reward, '浓度更高的地方奖励')
            # print(Gradient_reward, '梯度变化奖励')
            # print(Closer_reward, '靠近奖励')

            reward = k5 * exploration_bonus + \
                     k6 * dynamic_Probability + \
                     k2 * distance_reward + \
                     Gradient_reward + \
                     - 0.5  # 步数惩罚
            # reward = k5 * 0 + \
            #          k6 * dynamic_Probability + \
            #          k2 * 0 + \
            #          0 + \
            #          Closer_reward + \
            #          - 0.1  # 步数惩罚

            return reward

        self.element_collection[bot_name].ep_return += policy_()
        reward += policy_()

        if abs(next_this_point_mass - max_mass_thresh) < k4:
            self.element_collection[bot_name].ep_return = 200
            reward = 50
            done_s = True

            # taking_actions
        if self.element_collection[bot_name].stride_left < 1 or done_s:
            print('\n', "本轮得分:", '\n', self.element_collection[bot_name].ep_return)
            if view:
                plt.clf()
                plt.imshow(quality_scores_map, cmap='jet', vmin=0,
                           vmax=np.max(quality_scores_map))  # cmap 指定颜色映射
                plt.colorbar(label='Value')  # 添加颜色条
                plt.title('Heatmap Example')  # 设置标题
                plt.xlabel('X-axis')  # X 轴标签
                plt.ylabel('Y-axis')  # Y 轴标签
                plt.show()
            # plt.savefig(f"data/runing_record_map/{round(self.ep_return, 3)}.png")
            done_s = True  # 表明状态已经完成
        # 累计策略返回的分数
        # print(reward, 'reward')
        # print(self.element_collection[bot_name].ep_return, "ep_return\n")
        return np.stack([pos_grid_map, quality_scores_map]), reward, done_s

    # 设置控制模式

    def render(self, mode="parallel"):
        assert mode in ["parallel", "step"], "Invalid mode, must be either \"parallel\" or \"step\""
        # assert self.mass_fraction_map is not None
        if mode == "parallel":
            # 渲染的时候，直接开个新的线程进行并行渲染。
            t1 = threading.Thread(target=self.show, args=(), daemon=True)
            # 设置守护线程,即主线程启动才会启动子线程
            t1.start()
        elif mode == "step":
            pass

    def reset(self, bot_elements):
        # 重置每一个机器人的得分
        self.element_collection = self.create_bot(bot_elements)
        self.grid_map = np.zeros(self.grid_map.shape)  # 栅格概率地图重置


if __name__ == "__main__":
    # 创造一些机器人对象，便于获取
    Bot_elements = {"Bot1": BOT("Bot1", (0, 0), 200)}
    terminal = Terminal(Bot_elements, canvas_size=(900, 600), split_x_nums=10, split_y_nums=15)
    terminal.render()
    while True:
        terminal.draw_elements_on_canvas(csv_path=r'D:\Gas_detector\data\gas_mass_data\export8.csv')
        cv2.imshow("mass", cv2.resize(terminal.mass_fraction_map, (600, 900)))
        key = cv2.waitKey(60)
        for bot in terminal.element_collection.values():
            if key == 119:
                terminal.step(f"{bot.name}", 'upper')
            elif key == 115:
                terminal.step(f"{bot.name}", 'lower')
            elif key == 100:
                terminal.step(f"{bot.name}", 'right')
            elif key == 97:
                terminal.step(f"{bot.name}", 'left')
            elif key == 113:
                terminal.reset(Bot_elements)
                continue
        # terminal.draw_elements_on_canvas(csv_path=r'D:\Gas_detector\data\gas_mass_data\export7.csv')
        # print(terminal.element_collection['Bot2'].grid_pos_list)
        # print(terminal.element_collection['Bot2'].pixel_list)
        # terminal.step('Bot1', 'right')
        # terminal.step('Bot1', 'lower')
        # terminal.take_actions('lower', False, 'Bot2')
        # terminal.take_actions('right', False, 'Bot1')
        # cv2.waitKey(10)
